{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "46e9b6da-3780-47f3-a736-85086646472a",
   "metadata": {},
   "source": [
    "# Breaking Pipelined Hardware AES on CW305 FPGA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62ff3e73-f606-45f2-8781-7d4566627524",
   "metadata": {},
   "source": [
    "What happens when you take a hardware AES implementation and pipeline it?\n",
    "\n",
    "## Prerequisites\n",
    "- [PA_HW_CW305_1-Attacking_AES_on_an_FPGA.ipynb](PA_HW_CW305_1-Attacking_AES_on_an_FPGA.ipynb) (run this first!)\n",
    "- basic knowledge of Verilog\n",
    "- basic knowledge of AES\n",
    "\n",
    "## Requirements\n",
    "- CW305 or CW312 A35 target\n",
    "- CW-Lite, Pro, or Husky\n",
    "\n",
    "## Background\n",
    "Our [\"normal\" hardware AES implementation](https://github.com/newaetech/chipwhisperer/tree/develop/firmware/fpgas/aes) performs one round of AES in one clock cycle; the 10 rounds of AES-128 are done in 11 consecutive clock cycles (the +1 will be covered later). This is pretty fast! It means that 128 bits are encrypted in 11 cycles. At 100 MHz, this is 1.1 Gbps of encryption. For some applications -- think disk encryption, memory encryption, layer 2 or layer 3 encryption (MACsec/IPsec) -- this may not be fast enough.\n",
    "\n",
    "Some (but not all) [block cipher modes](https://en.wikipedia.org/wiki/Block_cipher_mode_of_operation) of AES can be pipelined: the idea is to replicate the round encryption for each round. So in the case of AES-128, imagine having 10 round modules; let's label these RM1 through RM10. In the first clock cycle, RM1 performs the first round of encryption on the first AES block. On the second clock cycle, the output of RM1 is fed to RM2, which performs the second round of encryption on the first AES block, **while at the same time** RM1 begins encrypting the second AES block. Once the pipeline is filled, ten different AES plaintext are at different stages of encryption, and a fully encrypted AES block is produced at each clock cycle. This boosts throughput by a factor of 11. Pipelining trades off area and power for throughput: a factor of X increase in throughput can be obtained at the cost of increasing area and power by X also. A block cipher mode can be pipelined if the encryption or decryption of a block does not depend on the result of the previous encryption or decryption; examples include CTR, XTS, and GCM. An example of a mode which cannot be pipelined is CBC.\n",
    "\n",
    "## The Target\n",
    "[The target(s)](https://github.com/newaetech/chipwhisperer/tree/develop/firmware/fpgas/aes_pipelined) for this demo were built quite straightforwardly by taking our [\"normal\" hardware AES implementation](https://github.com/newaetech/chipwhisperer/tree/develop/firmware/fpgas/aes), instantiating multiple rounds, and tweaking the control logic to feed the rounds. The goal was to minimize modifications to the original design; this lets us quantify the effect of pipelining on side-channel attacks. We'll study 4 pipelined variants: one fully pipelined implementation, and one \"half-pipelined\" implementation done three different ways.\n",
    "\n",
    "## Objectives\n",
    "1. See that encrypting multiple AES rounds at the same time does not prevent CPA attacks.\n",
    "2. De-mystify leakage models by matching different leakage models to different implementations.\n",
    "3. As a bonus: see that Vivado can do unexpected things that have a huge impact on leakage.\n",
    "\n",
    "## Platform Notes\n",
    "This notebook was developed with the CW305 target and CW-Husky. Other combinations are possible as per the requirements above, but results may differ slightly, and some tweaks will be required; these will be noted throughout (i.e. don't blindly run through the cells and expect it to work; follow the instructions and make any required modifications).\n",
    "\n",
    "In the case of the CW312 A35 target, the gain is set at a proper setting for the inductive shunt. If you're using a different shunt, you may need to make adjustments.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf27aa4-700c-438e-bf3f-59ed82e26102",
   "metadata": {},
   "outputs": [],
   "source": [
    "TARGET_PLATFORM = 'CW305_A100'\n",
    "#TARGET_PLATFORM = 'CW305_A35'\n",
    "#TARGET_PLATFORM = 'CW312T_A35'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d71522fd-0f8a-454c-aa36-4af37c10b032",
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer as cw\n",
    "scope = cw.scope()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5153da72-4de5-4c40-b038-552bffe41fd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "scope.default_setup()\n",
    "scope.io.hs2 = \"disabled\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07b0c60f-ebc0-428b-905d-a806d082d0e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "def program_target(half_pipe, target=None, force=True):\n",
    "    try:\n",
    "        if target is not None:\n",
    "            target.dis()        \n",
    "    except:\n",
    "        pass\n",
    "\n",
    "    if TARGET_PLATFORM == 'CW312T_A35':\n",
    "        scope.io.hs2 = 'clkgen'\n",
    "        fpga_id = 'cw312t_a35'\n",
    "        platform = 'ss2'\n",
    "    else:\n",
    "        platform = 'cw305'\n",
    "        if TARGET_PLATFORM == 'CW305_A100':\n",
    "            fpga_id = '100t'\n",
    "        elif TARGET_PLATFORM == 'CW305_A35':\n",
    "            fpga_id = '35t'\n",
    "\n",
    "    target = cw.target(scope, cw.targets.CW305_AES_PIPELINED, force=force, fpga_id=fpga_id, platform=platform, version=half_pipe)\n",
    "        \n",
    "    time.sleep(0.5)\n",
    "    scope.errors.clear()\n",
    "    assert target.half_pipe == half_pipe\n",
    "    return target"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbc3d8bf-6ad0-4d12-94c3-d8416c3a9aa8",
   "metadata": {},
   "source": [
    "# Fully Pipelined Target\n",
    "In this demo we'll look at a few different AES pipeline implementations.\n",
    "\n",
    "We'll start with a \"fully pipelined\" implementation, so named because it has dedicated hardware resources for each AES round; this lets it process one 128-bit AES block per clock cycle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8db756e4-b750-4606-bd12-5c8147686474",
   "metadata": {},
   "outputs": [],
   "source": [
    "target = program_target(half_pipe=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "381cd5b0-bd93-40a9-93a5-9f14f9aab8e2",
   "metadata": {},
   "source": [
    "Next we set up the clocks: CW-Husky requires a different setup when the ADC clock is driven by the target. If using the CW312T-A35 target, the capture hardware needs to drive the target clock."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adc9a39b-5e29-47b8-9153-147193f532f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "if 'CW305' in TARGET_PLATFORM:\n",
    "    scope.adc.offset = 0\n",
    "    target.vccint_set(1.0)\n",
    "    # we only need PLL1:\n",
    "    target.pll.pll_enable_set(True)\n",
    "    target.pll.pll_outenable_set(False, 0)\n",
    "    target.pll.pll_outenable_set(True, 1)\n",
    "    target.pll.pll_outenable_set(False, 2)\n",
    "\n",
    "    # run at 10 MHz:\n",
    "    target.pll.pll_outfreq_set(10E6, 1)\n",
    "\n",
    "    # 1ms is plenty of idling time\n",
    "    target.clkusbautooff = True\n",
    "    target.clksleeptime = 1\n",
    "    \n",
    "    if scope._is_husky:\n",
    "        scope.clock.clkgen_freq = 10e6\n",
    "        scope.clock.clkgen_src = 'extclk'\n",
    "        scope.clock.adc_mul = 4\n",
    "        # if the target PLL frequency is changed, the above must also be changed accordingly\n",
    "    else:\n",
    "        scope.clock.adc_src = \"extclk_x4\"\n",
    "\n",
    "else:\n",
    "    scope.adc.offset = 6\n",
    "    scope.clock.clkgen_freq = 7.37e6\n",
    "    scope.io.hs2 = 'clkgen'\n",
    "    if scope._is_husky:\n",
    "        scope.clock.clkgen_src = 'system'\n",
    "        scope.clock.adc_mul = 4\n",
    "        scope.clock.reset_dcms()\n",
    "    else:\n",
    "        scope.clock.adc_src = \"clkgen_x4\"\n",
    "    import time\n",
    "    time.sleep(0.1)\n",
    "    target._ss2_test_echo()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f58e8f27-4fd4-4f58-902d-e7ea405a0c2b",
   "metadata": {},
   "source": [
    "Finally, ensure the ADC clock is locked:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fc21743-f51e-4b2d-9a21-9e781004f21d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "for i in range(5):\n",
    "    scope.clock.reset_adc()\n",
    "    time.sleep(1)\n",
    "    if scope.clock.adc_locked:\n",
    "        break \n",
    "assert (scope.clock.adc_locked), \"ADC failed to lock\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8205d392-ff45-4bd6-8771-782672f70448",
   "metadata": {},
   "source": [
    "Occasionally the ADC will fail to lock on the first try; when that happens, the above assertion will fail (and on the CW-Lite, the red LED will be on). Simply re-running the above cell again should fix things."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebe3a283-ad22-42c2-a558-dcd860236051",
   "metadata": {},
   "source": [
    "## Trace Capture: one AES block at a time\n",
    "\n",
    "To simplify controlling the target's AES pipeline and how busy it's kept, the target has a FIFO  that's 128-bits wide and `target.fifo_depth` deep (on the CW305 this is 8192; on the CW312T_A35 it's 4096).\n",
    "Plaintexts to be encrypted are written to the FIFO; then when the target receives the \"go encrypt!\" signal, it processes everything that's in the FIFO as fast as it can (i.e. it will encrypt at a rate of one block per cycle until the FIFO is empty).\n",
    "\n",
    "To keep things simple as we get started, we'll start by having the target encrypt a single AES block at a time (e.g. write one block to the FIFO; make it \"go\"; capture the power trace; repeat).\n",
    "\n",
    "This is similar to how our [\"regular\" AES FPGA target](PA_HW_CW305_1-Attacking_AES_on_an_FPGA.ipynb) works, and it doesn't take advantage of this pipelined implementation's higher potential throughput, but it will let us learn some important things about the leakage that's present here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a03b22a4-5f49-4a30-8eaa-345b956112eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_file = \"projects/Tutorial_HW_CW305_AES_PIPELINED_HALF\" + str(target.half_pipe) + \".cwp\"\n",
    "project = cw.create_project(project_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6753efa-1f44-4006-ac56-2d225a57693e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tnrange\n",
    "import numpy as np\n",
    "import time\n",
    "from Crypto.Cipher import AES\n",
    "\n",
    "def get_traces(project, N, gain=None, check=True):\n",
    "    scope.adc.samples = 80\n",
    "    if gain == None:\n",
    "        if TARGET_PLATFORM == 'CW312T_A35':\n",
    "            scope.gain.db = 42\n",
    "        else:\n",
    "            scope.gain.db = 25\n",
    "    else:\n",
    "        scope.gain.db = gain\n",
    "    ktp = cw.ktp.Basic()\n",
    "    key, text = ktp.next()\n",
    "    # flush FIFO in case it's not empty:\n",
    "    target.fifo_flush()\n",
    "    for i in tnrange(N, desc='Capturing traces'):\n",
    "        key, text = ktp.next()  # manual creation of a key, text pair can be substituted here\n",
    "        ret = target.capture_trace(scope, text, key, pre_expand=False, check=check)\n",
    "        if not ret:\n",
    "            print(\"Failed capture\")\n",
    "            continue\n",
    "        project.traces.append(ret)\n",
    "        assert scope.adc.trig_count == 44\n",
    "    project.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef00c1f1-76c5-4127-86c1-1b69428c9d05",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_traces(project, 3000, check=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5850eb1-3871-4a8f-bec5-7ced868465a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.palettes import inferno\n",
    "from bokeh.plotting import figure, show\n",
    "from bokeh.resources import INLINE\n",
    "from bokeh.io import output_notebook\n",
    "from bokeh.models import Span, Legend, LegendItem\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d4a0041-66e0-4819-b01b-5ff93e3fa30a",
   "metadata": {},
   "outputs": [],
   "source": [
    "numtraces = 100\n",
    "output_notebook(INLINE)\n",
    "B = figure(width=1800)\n",
    "colors = itertools.cycle(inferno(numtraces))\n",
    "for i in range(numtraces):\n",
    "    B.line(list(range(scope.adc.samples)), project.traces[i].wave, color=next(colors))\n",
    "show(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b56b68ac-ce6a-475e-bb1f-c896a9a0f0ac",
   "metadata": {},
   "source": [
    "Traces look pretty similar to those of our regular AES FPGA target; you can clearly see the 10 AES rounds.\n",
    "\n",
    "When used this way, the only difference between the pipelined and non-pipelined AES targets is that in the non-pipelined target, the same logic gates are used for each round; here in the pipelined version, each round is a distinct and separate set of gates.\n",
    "\n",
    "Next let's look at what this means for side-channel attacks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "145ed5e8-30d3-4058-a879-e36c1ae80b88",
   "metadata": {},
   "source": [
    "## CPA Attack"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8183d88b-6428-4b74-9e1d-f9b87a7c8128",
   "metadata": {},
   "source": [
    "Let's see what we get from the same attack and leakage model as the [non-pipelined target](PA_HW_CW305_1-Attacking_AES_on_an_FPGA.ipynb):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91d8bf35-8976-4d3c-90e8-d06881f1893c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b97ff4f-f212-4f3f-8f08-ca9c74576a71",
   "metadata": {},
   "outputs": [],
   "source": [
    "attack = cwa.cpa(project, cwa.leakage_models.last_round_state_diff)\n",
    "attack.trace_range=[0,1000]\n",
    "cb = cwa.get_jupyter_callback(attack)\n",
    "attack_results = attack.run(cb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60356d4d-99b8-4815-a249-3f296b46da8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.average(attack_results.pge)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33097cd5-224a-4fc0-8911-bbd4f2ef82e7",
   "metadata": {},
   "source": [
    "Recall that the PGE for each key byte represents how far each correct key byte is from being the guessed key byte.\n",
    "\n",
    "The resulting average PGE of the above attack should be very close to 128, which is what you'd get if you randomly guessed the key.\n",
    "\n",
    "Further, if you watched the table evolve during the attack, you may have noticed that some key bytes get close to being correctly guessed but then get further again.\n",
    "\n",
    "These are all signs that we're not employing an appropriate or useful leakage model.\n",
    "\n",
    "If you follow the source code, you'll find that `cwa.leakage_models.last_round_state_diff` points to this leakage definition in [AES_128_8bit.py](https://github.com/newaetech/chipwhisperer/blob/develop/software/chipwhisperer/analyzer/attacks/models/AES128_8bit.py):\n",
    "```python\n",
    "def leakage(self, pt, ct, key, bnum):\n",
    "    # HD Leakage of AES State between 9th and 10th Round\n",
    "    # Used to break SASEBO-GII / SAKURA-G\n",
    "    st10 = ct[self.INVSHIFT_undo[bnum]]\n",
    "    st9 = inv_sbox(ct[bnum] ^ key[bnum])\n",
    "    return (st9 ^ st10)\n",
    "```\n",
    "\n",
    "(hot tip: in Jupyter, running `cwa.leakage_models.last_round_state.model??` will show you the source code for that function)\n",
    "\n",
    "Understanding why this leakage model works requires knowledge of the steps of AES and of the [AES implementation](https://github.com/newaetech/chipwhisperer/blob/develop/firmware/fpgas/cryptosrc/aes_googlevault/aes_core.v) and is not obvious on first glance, but at a high level it's looking at one byte of the AES `state` register in `aes_core.v` and calculates the XOR of <this byte after round 9 (st9)> and <this byte after round 10 (st10)>.\n",
    "\n",
    "Because of how AES works, we can do this on a byte-per-byte basis if we know the final ciphertext (`ct`); this is what allows the attack to guess one key byte at a time.\n",
    "\n",
    "Why then does this attack fail on the pipelined target? The Verilog code for the pipelined target is very similar to the non-pipelined target (compare [aes_core.v](https://github.com/newaetech/chipwhisperer/blob/develop/firmware/fpgas/cryptosrc/aes_googlevault/aes_core.v) and [aes_round.v](https://github.com/newaetech/chipwhisperer/blob/develop/firmware/fpgas/aes_pipelined/hdl/aes_round.v)), with one key difference: to enable pipelining, there is now a \"state\" register (`data_o` in `aes_round.v`) in *every pipeline stage*.\n",
    "\n",
    "Each of these state registers gets updated after processing a round; for each one of them, its content changes from what it held when it was used in processing the *previous plaintext* to what it holds now after processing the *current plaintext*. You can see how the core functions by running the very basic Verilog testbench provided [here](https://github.com/newaetech/chipwhisperer/tree/develop/firmware/fpgas/aes_pipelined/sim) with `make HALF_PIPE=0 DUMP=1`\n",
    "\n",
    "So, `cwa.leakage_models.last_round_state_diff` worked because there is a state register in `aes_core.v` which contains the output of round 9 and is then updated to the output of round 10, *for the same plaintext*. No such register exists in the pipelined `aes_round.v` because here, the output of each round goes to a register associated with only that round.\n",
    "\n",
    "But we are not out of luck! In fact the solution has already been outlined: we simply need to pick a round register and define a leakage function based on the state register content on the previous plaintext versus the current plaintext:\n",
    "\n",
    "```Python\n",
    "def leakage(self, pt, ct, prev_ct, key, bnum):\n",
    "    curr = inv_sbox(ct[bnum] ^ key[bnum])\n",
    "    prev = inv_sbox(prev_ct[bnum] ^ key[bnum])\n",
    "    return curr ^ prev\n",
    "```\n",
    "\n",
    "In the `leakage()` function above, `prev` represents the contents of the round 9 state register after it processed `prev_ct`, while `curr` represents the content of that same register when it was next updated, after it processed `ct`.\n",
    "\n",
    "Let's try it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13b2fbdd-4f0a-4b72-9f85-37d8d8a51957",
   "metadata": {},
   "outputs": [],
   "source": [
    "attack_pipe = cwa.cpa(project, cwa.leakage_models.pipeline_diff)\n",
    "attack_pipe.trace_range=[0,1000]\n",
    "cb = cwa.get_jupyter_callback(attack_pipe)\n",
    "attack_pipe_results = attack_pipe.run(cb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd2c41e8-0941-414a-a3b9-56106247d6c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.average(attack_pipe_results.pge)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe2877ce-1acf-4b7b-be5d-45c430b3e0b6",
   "metadata": {},
   "source": [
    "This isn't enough traces to retrieve the key, but we're quite close! You should see a PGE that's much closer to 0 than it is to 128 (with the CW312T_A35 you may need more traces to see this).\n",
    "\n",
    "Before we continue, let's define our own CPA attack function; it will run a bit faster because it processes all traces at once, instead of incrementally, and it lets us focus on the PGE, which will be handy as we will evaluate different leakage models.\n",
    "\n",
    "This uses the same approach as [Lab 4.2](<../courses/sca101/Lab 4_2 - CPA on Firmware Implementation of AES (MAIN).ipynb>):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a439b97f-f61a-45c7-867c-95926c41b6c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "HW = [bin(n).count(\"1\") for n in range(0, 256)]\n",
    "\n",
    "def mean(X):\n",
    "    return np.sum(X, axis=0)/len(X)\n",
    "\n",
    "def std_dev(X, X_bar):\n",
    "    return np.sqrt(np.sum((X-X_bar)**2, axis=0))\n",
    "\n",
    "def cov(X, X_bar, Y, Y_bar):\n",
    "    return np.sum((X-X_bar)*(Y-Y_bar), axis=0)\n",
    "\n",
    "def do_cpa(project, model, trace_range=None, point_range=None):\n",
    "    short_traces = []\n",
    "\n",
    "    if trace_range == None:\n",
    "        tstart = 0\n",
    "        tstop = len(project.traces)\n",
    "    else:\n",
    "        tstart = trace_range[0]\n",
    "        tstop = trace_range[1]\n",
    "\n",
    "    if point_range == None:\n",
    "        pstart = 0\n",
    "        pstop = len(project.traces[0].wave)\n",
    "    else:\n",
    "        pstart = point_range[0]\n",
    "        pstop = point_range[1]\n",
    "\n",
    "    # careful as this can lead to confusion! it means that ciphouts[i+1] is the ciphertext\n",
    "    # corresponding to short_traces[i] (and so ciphouts[i] is the corresponding *previous* ciphertext)\n",
    "    if tstart > 0:\n",
    "        ciphouts = [project.traces[tstart-1].textout]\n",
    "    else:\n",
    "        ciphouts = [[0]*16]\n",
    "\n",
    "    for i in range(tstart, tstop):\n",
    "        short_traces.append(project.traces[i].wave[pstart:pstop])\n",
    "        ciphouts.append(list(project.traces[i].textout))\n",
    "\n",
    "    num_traces = tstop - tstart\n",
    "\n",
    "    t_bar = np.sum(short_traces, axis=0)/num_traces\n",
    "    o_t = np.sqrt(np.sum((short_traces - t_bar)**2, axis=0))\n",
    "\n",
    "    cparefs = [0] * 16\n",
    "    bestguess = [0] * 16\n",
    "    bestguesses = []\n",
    "\n",
    "    for bnum in tnrange(0, 16):\n",
    "        maxcpa = [0] * 256\n",
    "        klist = [0]*16\n",
    "        for kguess in range(0, 256):\n",
    "            klist[bnum] = kguess\n",
    "            if model._has_prev:\n",
    "                hws = np.array([[HW[model.modelobj.leakage(None, ciphouts[i+1], None, ciphouts[i], klist, bnum)] for i in range(num_traces)]]).transpose()\n",
    "            else:\n",
    "                hws = np.array([[HW[model.modelobj.leakage(None, ciphouts[i+1], klist, bnum)] for i in range(num_traces)]]).transpose()\n",
    "            \n",
    "            hws_bar = mean(hws)\n",
    "            o_hws = std_dev(hws, hws_bar)\n",
    "            correlation = cov(short_traces, t_bar, hws, hws_bar)\n",
    "            cpaoutput = correlation/(o_t*o_hws)\n",
    "            maxcpa[kguess] = max(abs(cpaoutput))\n",
    "        bestguess[bnum] = np.argmax(maxcpa)\n",
    "        bestguesses.append(np.argsort(maxcpa)[::-1])\n",
    "        cparefs[bnum] = max(maxcpa)\n",
    "    \n",
    "    correct_recovered_key = model.modelobj.process_known_key(project.traces[0].key)\n",
    "    scores = []\n",
    "    for b in range(16):\n",
    "        score = list(bestguesses[b]).index(correct_recovered_key[b])\n",
    "        scores.append(score)\n",
    "    print('Remaining PGE: %f' % np.average(scores))\n",
    "\n",
    "    return bestguess, bestguesses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcc16fb0-b3b3-41f7-9b9e-b5fe0b19d7c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.pipeline_diff, point_range=[44,54])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38eececf-96fe-4fe2-842a-22d98836ab7a",
   "metadata": {},
   "source": [
    "The result should be 0 or very very close to it, which is much better than `attack_pipe_results` above because:\n",
    "1. We're using all 3000 traces.\n",
    "2. We're only looking at samples 44 to 54 from each power trace, which is where we know the leakage we're exploiting is occurring (because that's when the target round register gets updated).\n",
    "\n",
    "(Again, if you're using the CW312T_A35 target, you may find you need more traces.)\n",
    "\n",
    "Targeting the power samples more precisely will become more important for the next step, when we try the attack with the pipeline fully occupied.\n",
    "\n",
    "You can experiment with different settings for `point_range`; if you move it around a bit you should find slightly worse results; if you exclude samples `[44,54]` completely, you'll get much worse results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a53889be-fe09-463a-86c0-76e4af01e005",
   "metadata": {},
   "source": [
    "# Filling the Pipeline:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a304f48-ed62-4930-9f65-03e03cd75b07",
   "metadata": {},
   "source": [
    "We define a new trace capture function for this scenario.\n",
    "\n",
    "Recall that filling the pipeline involves writing all our plaintexts to the target's input FIFO.\n",
    "\n",
    "`target.capture_trace()` is a custom capture function for this target which takes care of this for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a7f54b5-c281-404a-9e5d-680e839cfea6",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_file = \"projects/Tutorial_HW_CW305_AES_PIPELINED_FILLED_HALF\" + str(target.half_pipe) + \".cwp\"\n",
    "project_filled = cw.create_project(project_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0592538-e6c0-4cb2-913f-272d3fd3705b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_traces_filled_pipeline(project, N, NPT, half_pipe=False, gain=None, check=True):\n",
    "    \"\"\" Args:\n",
    "        N: number of traces\n",
    "        NPT: number of encryptions per trace\n",
    "        half_pipe: set to True for half-pipelined targets, False for fully-pipelined targets\n",
    "        gain: leave to None to use defaults, otherwise provide desired gain in dB\n",
    "    \"\"\"\n",
    "    if half_pipe:\n",
    "        mx = 2\n",
    "    else:\n",
    "        mx = 1\n",
    "    scope.adc.samples = (NPT+100)*scope.clock.adc_mul*mx\n",
    "    if gain == None:\n",
    "        if TARGET_PLATFORM == 'CW312T_A35':\n",
    "            scope.gain.db = 34\n",
    "        else:\n",
    "            scope.gain.db = 18\n",
    "    else:\n",
    "        scope.gain.db = gain\n",
    "    ktp = cw.ktp.Basic()\n",
    "    key, text = ktp.next()\n",
    "    # flush FIFO in case it's not empty:\n",
    "    target.fifo_flush()\n",
    "    for i in tnrange(N, desc='Capturing traces'):\n",
    "        texts = []\n",
    "        for j in range(NPT):\n",
    "            key, text = ktp.next()\n",
    "            texts.append(list(text))\n",
    "        ret = target.capture_trace(scope, texts, key, pre_expand=False, check=check)\n",
    "        if not ret:\n",
    "            print(\"Failed capture\")\n",
    "            continue\n",
    "        project.traces.append(ret)\n",
    "    project.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a4c6b11-ae66-47e9-838a-7863fb2cc462",
   "metadata": {},
   "source": [
    "`get_traces_filled_pipeline()` has a new argument, `NPT`.\n",
    "\n",
    "This is the number of input words that are written to the target FIFO before it is made to go. This can be repeated `N` times.\n",
    "\n",
    "So, `get_traces_filled_pipeline()` collects `N` power traces, each of which contains `NPT` encryptions.\n",
    "\n",
    "`NPT` could be anything, but to get interesting results it should be greater than the number of AES rounds (so that at least part of the power trace has the pipeline fully occupied).\n",
    "\n",
    "For fun let's see what traces with `NPT=20` looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ece5b72f-ae57-4701-96fd-b31348d485bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_traces_filled_pipeline(project_filled, N=10, NPT=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5741188e-9b66-4be7-8c5b-51bcc7298817",
   "metadata": {},
   "outputs": [],
   "source": [
    "numtraces = len(project_filled.traces)\n",
    "output_notebook(INLINE)\n",
    "B = figure(width=1800)\n",
    "colors = itertools.cycle(inferno(numtraces))\n",
    "for i in range(numtraces):\n",
    "    B.line(list(range(scope.adc.samples)), project_filled.traces[i].wave, color=next(colors))\n",
    "show(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbeb2e85-be40-4649-8ba1-9245f1c0be6c",
   "metadata": {},
   "source": [
    "You can clearly see power consumption ramping up with each clock cycle as the AES pipeline fills at the start, then go back down as it empties at the end.\n",
    "\n",
    "Since we're interested in the feasibility of a side-channel attack on an AES pipeline, let's maximize `NPT`.\n",
    "\n",
    "Notes:\n",
    "1. Each iteration involves sending `target.fifo_depth` plaintexts and retrieving the same number of ciphertexts, so it takes a while; don't be concerned if you don't see any progress for a while.\n",
    "2. If you have a CW-lite, you'll have to reduce NPT to 6000 so that you don't exceed the Lite's maximum capture size; you can increase N to 4 to capture roughly the same total number of encryptions.\n",
    "\n",
    "`setN()` is a convenience function which adjusts N for different targets, in consideration for their differing FIFO depths, and the fact that the CW312 target requires more traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bd90924-6360-41d3-9338-d020b87008b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def setN(N):\n",
    "    if 'A35' in TARGET_PLATFORM:\n",
    "        N=N*2 # double N because `target.fifo_depth` is halved on A35 targets; double it again because more traces are required on this target\n",
    "    if 'CW312' in TARGET_PLATFORM:\n",
    "        N=N*2 # double N because more traces are required on this target\n",
    "    return N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "883aa960-1339-4632-89ce-94660580bcf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_filled = cw.create_project(project_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e98f43eb-fba7-4599-84f1-2d92c3d19bb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = setN(3)\n",
    "get_traces_filled_pipeline(project_filled, N=N, NPT=target.fifo_depth, check=False) # adjust N, NPT as per notes above if necessary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db000979-4b9d-4943-b9c3-a745dcaec388",
   "metadata": {},
   "source": [
    "Now we have `N` long power traces, each containing `NPT` encryptions.\n",
    "\n",
    "Our CPA attack needs to associate one power trace to each encryption. How do we do that? By splitting up each power trace into `NPT` segments! `split_traces()` is a convenience function that does this for us.\n",
    "\n",
    "This is where the `point_range` from before comes in handy, because we need to tell `split_traces()` *where* to split up the power trace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef2ae81a-687a-4ffa-b8d4-92bf49f33423",
   "metadata": {},
   "outputs": [],
   "source": [
    "split_traces = target.split_traces(scope, project_filled.traces, 44, 54)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f40aede5-5367-43a3-bc26-5f9832228f78",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_file = \"projects/Tutorial_HW_CW305_AES_PIPELINED_FILLED_SPLIT_HALF\" + str(target.half_pipe) + \".cwp\"\n",
    "project_split = cw.create_project(project_file, overwrite=True)\n",
    "\n",
    "for t in split_traces:\n",
    "    project_split.traces.append(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ee19e28-9661-4b5c-a8dc-82e0232e6a12",
   "metadata": {},
   "source": [
    "Now we have our synthesized per-encryption traces in this new project.\n",
    "\n",
    "`split_traces()` has also done the work of associating the plaintext and ciphertext that goes which each power trace segment, so `project_split` is ready to be fed to `do_cpa()`.\n",
    "\n",
    "(note that we no longer provide the `point_range` argument because it's already been used to build our set of power trace segments)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f1c45c4-5974-424f-9c34-1feb94bf912a",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project_split, cwa.leakage_models.pipeline_diff, trace_range=[0, 2000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71824b53-5825-4577-bfe8-fc95ac7300c2",
   "metadata": {},
   "source": [
    "2000 traces worked were sufficient (or almost) when the pipeline was processing only one encryption at a time; now that it's fully loaded, the noise of all the active pipeline stages means we need more traces\n",
    "The attack needs more traces to succeed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c54044c-40a3-4a21-9aad-c5a8685511cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project_split, cwa.leakage_models.pipeline_diff)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "609cc4f4-1ba6-41b9-9fc9-7639fea4e5bb",
   "metadata": {},
   "source": [
    "This should work! This used `N*(NPT-20)` traces (the first and last 10 segments of each large power trace are omitted because the pipeline isn't fully loaded for those).\n",
    "\n",
    "You can see just how many traces are required; it should be around 20000:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d20ba4a1-6558-445f-90a9-5cc332a3413e",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project_split, cwa.leakage_models.pipeline_diff, trace_range=[0, 20000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15b6a2e9-b182-486b-9433-4374ee87a1d4",
   "metadata": {},
   "source": [
    "So a ~10-fold increase in the noise that's present in the power traces (a very rough and simplistic approximation) leads to a ~10-fold increase in the number of traces required for the CPA attack.\n",
    "\n",
    "(On the CW312T-A35, around 50K traces should do it.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33bcc026-c999-4d87-9009-081889f3994d",
   "metadata": {},
   "source": [
    "# Different Ways to Pipeline\n",
    "\n",
    "Encryption is expensive; a 10-core AES pipeline is going to be about 10 times larger than the equivalent non-pipelined implementation, and if 128 bits/cycle throughput doesn't hit the sweet spot for a particular application, other landing points are possible.\n",
    "\n",
    "You can go faster by implementing multiple pipelines in parallel; you can also go slower by reducing the number of pipeline stages.\n",
    "\n",
    "Consider for example a five-stage pipeline, where each stage does two AES rounds (in two clock cycles) instead of just one. This will be half the size and half the throughput, since it will be able to take in one 16-byte plaintext every two cycles.\n",
    "\n",
    "Changing the number of pipeline stages from ten to five impacts the leakage model (if you understood the leakage model for the ten-stage implementation you should immediately see why); it will also be interesting to see how the reduced number of stages affects the number of traces required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "137a8fb4-eb89-4cae-97e9-f924f0830566",
   "metadata": {},
   "outputs": [],
   "source": [
    "target.dis()\n",
    "target = program_target(half_pipe=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a065da93-e82f-42f9-ae04-dcddd44a853f",
   "metadata": {},
   "source": [
    "We start as before, with a single encryption at a time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfba3760-1252-4be7-a164-4dc2a9b3137e",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_file = \"projects/Tutorial_HW_CW305_AES_PIPELINED_HALF\" + str(target.half_pipe) + \".cwp\"\n",
    "project = cw.create_project(project_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "810764c0-f8ba-490b-b121-fe66b3c38413",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_traces(project, 2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a61b40-b2c4-47f0-a257-55e17be79d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "project.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fc621f2-4e90-487b-94c1-42763f13075b",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.pipeline_diff, point_range=[44,54])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fd09562-111c-44e2-9ca8-b4160d1916fd",
   "metadata": {},
   "source": [
    "Yep, the leakage model that worked for the fully-pipelined target is useless here!\n",
    "\n",
    "To figure out what the leakage model should be, we again need to look at the implementation details.\n",
    "\n",
    "In `aes_core.v`, the initial round key addition is essentially done as a distinct \"round 0\", in its own clock cycle.\n",
    "\n",
    "This means that the full encrytion is done in 11 cycles (not 10), so when we halve the pipeline, we have two implementation choices:\n",
    "\n",
    "|pipe stage |version 1 rounds |version 2 rounds |\n",
    "|-----------|----------|----------|\n",
    "| 1         | 0        | 0, 1     |\n",
    "| 2         | 1, 2     | 2, 3     |\n",
    "| 3         | 3, 4     | 4, 5     |\n",
    "| 4         | 5, 6     | 6, 7     |\n",
    "| 5         | 7, 8     | 8, 9     |\n",
    "| 6         | 9, 10    | 10       |\n",
    "\n",
    "Our current target uses version 1. Again you may find it helpful to run the tesbench that's provided [here](https://github.com/newaetech/chipwhisperer/tree/develop/firmware/fpgas/aes_pipelined/sim) to see the core in action: `make HALF_PIPE=1 DUMP=1`.\n",
    "\n",
    "One thing you should see right away is that since the last stage processes rounds 9 and 10 for the same plaintext in sequence, the leakage model that we use for the non-pipelined version of this target should work here!\n",
    "\n",
    "And it does:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14bae6a3-68da-461b-917d-dc7831b7dd36",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.last_round_state_diff, point_range=[48,58])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2f8e4dc-bcf9-4222-96aa-b00a779bb657",
   "metadata": {},
   "source": [
    "This targeted the pipeline stage 6 register as it goes from the round 9 output to the round 10 output (on the same plaintext).\n",
    "\n",
    "But we can also target that same register as it goes from the round 10 output for the *previous* plaintext to the round 9 output for the *current* plaintext, using this leakage model:\n",
    "\n",
    "```Python\n",
    "def leakage(self, pt, ct, prev_ct, key, bnum):\n",
    "    curr = inv_sbox(ct[bnum] ^ key[bnum])\n",
    "    prev = prev_ct[self.INVSHIFT_undo[bnum]]\n",
    "    return curr ^ prev\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71458c1b-4e6e-4d62-9131-86ff610e7c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.half_pipeline_diff, point_range=[48,58])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e893372-c4c2-4c53-8e95-691774257753",
   "metadata": {},
   "source": [
    "This also works, just not quite as well as `last_round_state_diff` (on the CW312T_A35 it seems to barely work; however if you use more traces you should see better results). Let's move on to fully loading the pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f3284c0-98e4-4e4c-b6cb-287d373b541e",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_file = \"projects/Tutorial_HW_CW305_AES_PIPELINED_FILLED_HALF\" + str(target.half_pipe) + \".cwp\"\n",
    "project_filled = cw.create_project(project_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f61c567-0f8a-46a2-968e-07ad369e2ed2",
   "metadata": {},
   "outputs": [],
   "source": [
    "N=setN(3)\n",
    "get_traces_filled_pipeline(project_filled, N=N, NPT=target.fifo_depth, half_pipe=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a05c82af-37a4-4c97-a5ee-a1d572a04cce",
   "metadata": {},
   "outputs": [],
   "source": [
    "split_traces = target.split_traces(scope, project_filled.traces, 48, 58)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bdcd405-49a5-4425-bae4-528af9a5c4b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_file = \"projects/Tutorial_HW_CW305_AES_PIPELINED_FILLED_SPLIT_HALF\" + str(target.half_pipe) + \".cwp\"\n",
    "project_split = cw.create_project(project_file, overwrite=True)\n",
    "\n",
    "for t in split_traces:\n",
    "    project_split.traces.append(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18602f9d-2fee-4a25-b711-3e4df5c17b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project_split, cwa.leakage_models.last_round_state_diff, trace_range=[0,10000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b265bea-e16a-4e75-8cce-48dc28919ca7",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project_split, cwa.leakage_models.half_pipeline_diff, trace_range=[0,10000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87301c81-680c-4f35-84a1-01d0eaa98eaa",
   "metadata": {},
   "source": [
    "`last_round_state_diff` still outperforms, but the delta is much smaller.\n",
    "\n",
    "What if we *combine* the two models?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d1dba4-6536-4911-a2cf-7645eb02f268",
   "metadata": {},
   "outputs": [],
   "source": [
    "def do_cpa2x(project, model1, model2, trace_range=None, point_range=None):\n",
    "    short_traces = []\n",
    "\n",
    "    if trace_range == None:\n",
    "        tstart = 0\n",
    "        tstop = len(project.traces)\n",
    "    else:\n",
    "        tstart = trace_range[0]\n",
    "        tstop = trace_range[1]\n",
    "\n",
    "    if point_range == None:\n",
    "        pstart = 0\n",
    "        pstop = len(project.traces[0].wave)\n",
    "    else:\n",
    "        pstart = point_range[0]\n",
    "        pstop = point_range[1]\n",
    "\n",
    "    # careful as this can lead to confusion! it means that ciphouts[i+1] is the ciphertext\n",
    "    # corresponding to short_traces[i] (and so ciphouts[i] is the corresponding *previous* ciphertext)\n",
    "    if tstart > 0:\n",
    "        ciphouts = [project.traces[tstart-1].textout]\n",
    "    else:\n",
    "        ciphouts = [[0]*16]\n",
    "\n",
    "    for i in range(tstart, tstop):\n",
    "        short_traces.append(project.traces[i].wave[pstart:pstop])\n",
    "        ciphouts.append(list(project.traces[i].textout))\n",
    "\n",
    "    num_traces = tstop - tstart\n",
    "\n",
    "    t_bar = np.sum(short_traces, axis=0)/num_traces\n",
    "    o_t = np.sqrt(np.sum((short_traces - t_bar)**2, axis=0))\n",
    "\n",
    "    cparefs = [0] * 16\n",
    "    bestguess = [0] * 16\n",
    "    bestguesses = []\n",
    "\n",
    "    for bnum in tnrange(0, 16):\n",
    "        maxcpa = [0] * 256\n",
    "        klist = [0]*16\n",
    "        for kguess in range(0, 256):\n",
    "            klist[bnum] = kguess\n",
    "            hwss = []\n",
    "            for model in [model1, model2]:\n",
    "                if model._has_prev:\n",
    "                    hwss.append(np.array([[HW[model.modelobj.leakage(None, ciphouts[i+1], None, ciphouts[i], klist, bnum)] for i in range(num_traces)]]).transpose())\n",
    "                else:\n",
    "                    hwss.append(np.array([[HW[model.modelobj.leakage(None, ciphouts[i+1], klist, bnum)] for i in range(num_traces)]]).transpose())\n",
    "            \n",
    "            for hws in hwss:\n",
    "                hws_bar = mean(hws)\n",
    "                o_hws = std_dev(hws, hws_bar)\n",
    "                correlation = cov(short_traces, t_bar, hws, hws_bar)\n",
    "                cpaoutput = correlation/(o_t*o_hws)\n",
    "                maxcpa[kguess] += max(abs(cpaoutput))        \n",
    "\n",
    "        bestguess[bnum] = np.argmax(maxcpa)\n",
    "        bestguesses.append(np.argsort(maxcpa)[::-1])\n",
    "        cparefs[bnum] = max(maxcpa)        \n",
    "\n",
    "    correct_recovered_key = model1.modelobj.process_known_key(project.traces[0].key)\n",
    "    scores = []\n",
    "    for b in range(16):\n",
    "        score = list(bestguesses[b]).index(correct_recovered_key[b])\n",
    "        scores.append(score)\n",
    "    print('Remaining PGE: %f' % np.average(scores))\n",
    "\n",
    "    return bestguess, bestguesses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e9f1477-3bda-4cbb-8cca-741c87e37e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa2x(project_split, cwa.leakage_models.half_pipeline_diff, cwa.leakage_models.last_round_state_diff, trace_range=[0,10000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08a3889b-c5fa-4923-a865-70582f25f872",
   "metadata": {},
   "source": [
    "It looks like the combination is beneficial! (Except on the CW312T_A35, where `last_round_state_diff` is so much better than `half_pipeline_diff`, and so their combination is not beneficial.)\n",
    "\n",
    "The number of traces needs to be bumped up to around 18000 in order to fully succeed. This is a modest reduction compared to the fully-pipeline implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1bf3792-c5b6-44a6-86ce-fec747f91206",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa2x(project_split, cwa.leakage_models.half_pipeline_diff, cwa.leakage_models.last_round_state_diff, trace_range=[0,18000])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2e4a6b8-fe26-4b46-b9ff-8ce19fdd9d54",
   "metadata": {},
   "source": [
    "# Half-Pipeline Version 2\n",
    "\n",
    "What if we use the second way to half-pipeline? The state registers are now updated as per the version 2 column of this table:\n",
    "\n",
    "|pipe stage |version 1 rounds |version 2 rounds |\n",
    "|-----------|----------|----------|\n",
    "| 1         | 0        | 0, 1     |\n",
    "| 2         | 1, 2     | 2, 3     |\n",
    "| 3         | 3, 4     | 4, 5     |\n",
    "| 4         | 5, 6     | 6, 7     |\n",
    "| 5         | 7, 8     | 8, 9     |\n",
    "| 6         | 9, 10    | 10       |\n",
    "\n",
    "What leakage can we exploit here?\n",
    "\n",
    "- The stage 6 register gets updated with the final ciphertext values for each encryption. If you code the leakage model for this, you'll find that the key is XOR's twice, which takes it out of the leakage definition; this can't be exploited to recover the key.\n",
    "- The stage 5 register first gets updated from the round 9 output of the previous ciphertext to the round 8 output of the next ciphertext. This is not easily exploitable because the round key touches multiple bytes between round 8 and round 9.\n",
    "- The stage 5 register then gets updated from the round 8 output to the round 9 output. Again, this is not easily exploitable because the round key touches multiple bytes between round 8 and round 9.\n",
    "\n",
    "What about the first pipeline stage? We've overlooked it until now. The \"round 0\" content is the initial round key addition. This can't be exploited across successive encryptions because the key would cancel out:\n",
    "\n",
    "```Python\n",
    "prev = prev_pt[bnum] ^ key[bnum]\n",
    "curr = pt[bnum] ^ key[bnum]\n",
    "curr ^ prev == pt[bnum] ^ prev_pt[bnum]\n",
    "```\n",
    "\n",
    "And the round 1 output has gone through the MixColumns step which diffuses the key across multiple bytes.\n",
    "\n",
    "Does that mean version 2 can't succumb to this CPA attack? Let's find out..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bede3fd2-e701-465e-9ed8-2ff077cefda2",
   "metadata": {},
   "outputs": [],
   "source": [
    "target.dis()\n",
    "target = program_target(half_pipe=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb3cf7a2-3dcb-49a9-8170-e7a496767ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_file = \"projects/Tutorial_HW_CW305_AES_PIPELINED_HALF\" + str(target.half_pipe) + \".cwp\"\n",
    "project = cw.create_project(project_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93955841-9a07-4cbe-bc7c-f1cfacfe8e7b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if 'CW305' in TARGET_PLATFORM:\n",
    "    gain = 22\n",
    "else:\n",
    "    gain = 41\n",
    "\n",
    "get_traces(project, 3000, gain=gain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31d0e043-a6c8-49e5-b1b9-b1dc511b22c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "project.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01f9d54b-474b-43ce-91c6-d0a67955f23d",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.half_pipeline_diff, point_range=[44,54])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de4c3ed1-3601-4c13-b7e4-98df39722ed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.last_round_state_diff, point_range=[44,54])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a83a705-952e-419d-a649-82591b6b6453",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.pipeline_diff, point_range=[44,54])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7411440-5500-4b05-ad05-49207dcff245",
   "metadata": {},
   "source": [
    "Surprisingly, two of the three leakage models that we've exploited for the previous targets work **very** well. **How can this be?!?**\n",
    "\n",
    "Recall that `half_pipeline_diff` was developed for version 1, where the last stage register was updated from the round 10 value of the previous ciphertext to the round 9 value of the current ciphertext.\n",
    "\n",
    "Here with version 2, the round 10 and round 9 values don't even go to the same physical state register.\n",
    "\n",
    "**So why does it work here?** (Not only does it work, it appears to work at least as well is it does against version 1!)\n",
    "\n",
    "To understand what seems to have happened, we'll have to dive into the Verilog and the actual FPGA implementation.\n",
    "\n",
    "Looking at the AES round instantiations in [`aes_half_pipeline_top.v`](https://github.com/newaetech/chipwhisperer/blob/develop/firmware/fpgas/aes_pipelined/hdl/aes_half_pipeline_top.v), when `HALF_PIPE=2` we have 6 instances of `aes_two_rounds`; the last instance of `aes_two_rounds` (k = 10) is used to perform only one round (the last round).\n",
    "\n",
    "This is lazy coding: the \"correct\" thing to do would be to instantiate `aes_round` separately for the last round, outside of the generate block, like is done for the initial round in the `HALF_PIPE=1` case. This really shouldn't matter -- the two cases should be logically equivalent -- but it appears Vivado had other ideas and did some unexpected \"optimizations\". The easiest way to see that something is amiss is to show the connectivity of the different round modules in Vivado's graphical device view. In the images below, the input FIFO logic is highlighted green and the output FIFO logic is highlighted red.\n",
    "\n",
    "On the left-hand side, the white lines show the connectivity from the last round (`gen_half_rounds[10].U_aes_two_rounds`); as expected, there are many connections to the output FIFOs.\n",
    "\n",
    "On the right-hand side, the white lines show the connectivity from the second last round (`gen_half_rounds[8].U_aes_two_rounds`); what is *very* surprising is that is that this *also* has many connections to the output FIFOs!\n",
    "\n",
    "<img src=\"img/aes_pipe_half2_last_round_connections.png\" width=\"600\"> <img src=\"img/aes_pipe_half2_2ndlast_round_connections.png\" width=\"600\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b97be2e-0457-4570-b03f-20c81498cb6f",
   "metadata": {},
   "source": [
    "(these are from the CW305 100t target; the other targets will look different but should show the same idea)\n",
    "\n",
    "Logically, this should not be happening: `gen_half_rounds[8].U_aes_two_rounds` should only be connected to the previous and next round instances; it should have no direct connection to the output FIFO.\n",
    "\n",
    "This can be verified in the [synthesized netlist](https://github.com/newaetech/chipwhisperer/blob/develop/firmware/fpgas/aes_pipelined/cw305_aes_pipelined_half2_netlist.v). What we're looking for is an output of `gen_half_rounds[10].U_aes_two_rounds` connected to an input of `gen_half_rounds[8].U_aes_two_rounds`.\n",
    "\n",
    "Unfortunately, Vivado mangles net names horribly so this is not as easy as it should be, but with a little bit of sweat and tears you can find that, for example, `gen_half_rounds[10].U_aes_two_rounds_n_39` is an example of such a signal.\n",
    "\n",
    "A shortcut (which unfortunately involves Tcl 😱), is to open the synthesized design in Vivado (for this you'll have to re-run the synthesis yourself) and query Vivado for outputs of `gen_half_rounds[10].U_aes_two_rounds` connected to inputs of `gen_half_rounds[8].U_aes_two_rounds`:\n",
    "```tcl\n",
    "set round10outs [get_nets -of [get_pins gen_half_pipe.U_aes_pipeline/gen_half_rounds[10].U_aes_two_rounds/* -filter {DIRECTION == OUT}]]\n",
    "set round8ins [get_nets -of [get_pins gen_half_pipe.U_aes_pipeline/gen_half_rounds[8].U_aes_two_rounds/* -filter {DIRECTION == IN}]]\n",
    "set intersect [list]\n",
    "foreach elem $round8ins {if {$elem in $round10outs} {lappend intersect $elem}}\n",
    "puts $intersect\n",
    "```\n",
    "\n",
    "(On the CW312T_A35 target, pre-pend the first argument of each `get_pins` call with `U_cw305_dut/`.)\n",
    "\n",
    "In our implementation we find 134 such signals."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6abe44f2-e08d-4ee7-adba-9210daad7a68",
   "metadata": {},
   "source": [
    "Let's confirm our hunch. Setting `HALF_PIPE` to 3 instantiates  [`aes_half_pipeline_top_fixed.v`](https://github.com/newaetech/chipwhisperer/blob/develop/firmware/fpgas/aes_pipelined/hdl/aes_half_pipeline_top_fixed.v), which is identical to `aes_half_pipeline_top.v` except that the last round is done by an instance of `aes_round` instead of `aes_two_rounds`.\n",
    "\n",
    "Before trying to attack this implementation, let's do a quick graphical sanity check (second-last round connections on the left; last round connections on the right):\n",
    "\n",
    "<img src=\"img/aes_pipe_half3_last_round_connections.png\" width=\"600\"> <img src=\"img/aes_pipe_half3_2ndlast_round_connections.png\" width=\"600\">\n",
    "\n",
    "This looks much better, and if you run the synthesis and the above Tcl commands, you should find that `$intersect` is now empty.\n",
    "\n",
    "Let's now quickly try the attack with all three models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2792a1b0-c1b2-4ff4-8110-659b176df20e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "target.dis()\n",
    "target = program_target(half_pipe=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bf26bb2-1d1f-442e-8e73-228f4cb70d02",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_file = \"projects/Tutorial_HW_CW305_AES_PIPELINED_HALF\" + str(target.half_pipe) + \".cwp\"\n",
    "project = cw.create_project(project_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "639e9e69-5eec-49ab-8537-af02801c4652",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_traces(project, 3000, gain=gain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56eac65c-0b7e-4a3e-bf2b-6feb20e28608",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.half_pipeline_diff, point_range=[44,54])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ba575b4-9da7-4015-82f3-1de825581b7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.last_round_state_diff, point_range=[44,54])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42a04a89-23ed-443e-b9cd-e0749dadd5dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = do_cpa(project, cwa.leakage_models.pipeline_diff, point_range=[44,54])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a11cd027-590f-4b69-8271-572b013a7320",
   "metadata": {},
   "source": [
    "**Success!**\n",
    "\n",
    "This doesn't mean that this implementation isn't vulnerable to CPA attacks (meet the [MixColumn attack](../courses/sca201/Lab%202_3%20-%20Attacking%20Across%20MixColumns.ipynb)!).\n",
    "\n",
    "The objective here was to illustrate the impact of design choices, both intentional and not, on leakage. By luck, we stumbled into a wonderful illustration of how implementation tools can do strange and unexpected things with our source code that alter its side-channel leakage, and that this behaviour is not exclusive to software compilers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "093a8921-73a6-4947-b6fe-1d3e5c53a47b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
